import cv2
import numpy as np
from PyQt5.QtWidgets import QMessageBox
from PyQt5.QtCore import Qt
from PyQt5 import QtWidgets
from segment_anything import SamPredictor, sam_model_registry


def init_sam_model(self):
    """初始化SAM模型"""
    if self.source_image is None:
        QMessageBox.warning(self, "警告", "请先打开图像")
        return False

    try:
        # 检查是否已经初始化了SAM
        if not hasattr(self, 'sam_predictor'):
            self.statusbar.showMessage("正在加载SAM模型，请稍等...")

            # 初始化SAM模型
            model_path = "sam_vit_l_0b3195.pth"  # 使用标准 SAM 模型
            sam = sam_model_registry["vit_l"](checkpoint=model_path)

            self.sam_predictor = SamPredictor(sam)

            # 将图像设置到SAM中（确保是RGB格式）
            self.sam_rgb_image = cv2.cvtColor(self.source_image, cv2.COLOR_BGR2RGB)
            self.sam_predictor.set_image(self.sam_rgb_image)

            self.statusbar.showMessage("SAM模型已加载，可以开始选点")
        else:
            # 如果已经初始化过但是图像变化了，需要重新设置图像
            self.sam_rgb_image = cv2.cvtColor(self.source_image, cv2.COLOR_BGR2RGB)
            self.sam_predictor.set_image(self.sam_rgb_image)

        return True
    except Exception as e:
        QMessageBox.critical(self, "错误", f"加载SAM模型失败: {str(e)}")
        self.statusbar.showMessage("SAM模型加载失败")
        return False


def btnSAM_Clicked(self):
    """初始化SAM模型并启用选点功能"""
    if not init_sam_model(self):
        return

    # 初始化选点列表和历史记录
    self.input_points = []
    self.input_labels = []
    self.input_points_history = []
    self.input_labels_history = []
    self.drawing_image_history = []

    # 制作图像副本以用于绘制点
    if self.source_image is not None:
        self.drawing_image = self.source_image.copy()

    # 启用选点模式
    self.sam_mode = True

    # 提示用户
    self.statusbar.showMessage("请在图像上左键点击添加前景点，右键点击添加背景点，按回车键进行分割，Ctrl+Z撤销选点")
    QMessageBox.information(self, "提示",
                            "左键：添加前景点(绿色)\n右键：添加背景点(红色)\n回车：执行分割\nCtrl+Z：撤销上一个选点")


def sam_add_point(self, x, y, is_positive=True):
    """添加选点到SAM模型"""
    if not self.sam_mode or self.source_image is None:
        return

    # 保存当前状态用于撤销
    self.input_points_history.append(self.input_points.copy())
    self.input_labels_history.append(self.input_labels.copy())
    if hasattr(self, 'drawing_image') and self.drawing_image is not None:
        self.drawing_image_history.append(self.drawing_image.copy())
    else:
        self.drawing_image_history.append(self.source_image.copy())

    # 计算原图上的坐标
    orig_x = int(x / self.scale_ratio)
    orig_y = int(y / self.scale_ratio)

    # 确保坐标不超出图像范围
    h, w = self.source_image.shape[:2]
    orig_x = max(0, min(orig_x, w - 1))
    orig_y = max(0, min(orig_y, h - 1))

    # 添加点到列表
    self.input_points.append([orig_x, orig_y])
    self.input_labels.append(1 if is_positive else 0)  # 1 为前景，0 为背景

    # 在图像副本上绘制点
    if not hasattr(self, 'drawing_image') or self.drawing_image is None:
        self.drawing_image = self.source_image.copy()

    color = (0, 255, 0) if is_positive else (0, 0, 255)  # 绿色为前景，红色为背景
    cv2.circle(self.drawing_image, (orig_x, orig_y), 5, color, -1)

    # 更新显示
    self.display_image(self.drawing_image)

    # 更新状态栏
    self.statusbar.showMessage(f"添加 {'前景' if is_positive else '背景'} 点: ({orig_x}, {orig_y})")


def sam_undo_point(self):
    """撤销最后添加的点"""
    if not self.sam_mode or not self.input_points_history:
        return

    # 恢复上一个状态
    self.input_points = self.input_points_history.pop()
    self.input_labels = self.input_labels_history.pop()
    self.drawing_image = self.drawing_image_history.pop()

    # 强制重新计算和显示图像，确保适应当前窗口大小
    QtWidgets.QApplication.processEvents()
    self.display_image(self.drawing_image)

    # 更新状态栏
    if self.input_points:
        self.statusbar.showMessage(f"已撤销最后一个选点，当前剩余 {len(self.input_points)} 个点")
    else:
        self.statusbar.showMessage("已撤销所有选点")


def sam_segment(self):
    """使用SAM进行分割"""
    if not self.input_points:
        QMessageBox.warning(self, "警告", "请先添加点再进行分割")
        return

    try:
        # 确保输入点的格式正确
        input_points = np.array(self.input_points)
        input_labels = np.array(self.input_labels)

        # 使用SAM进行预测
        masks, scores, logits = self.sam_predictor.predict(
            point_coords=input_points,
            point_labels=input_labels,
            multimask_output=False,  # 仅获取一个mask
        )

        # 获取分割后的mask
        mask = masks[0]
        self.mask = mask  # 保存掩码供窗口调整时使用

        # 将mask应用到原图上创建分割图像
        segmented_image = self.source_image.copy()
        # 创建半透明覆盖效果
        overlay = self.source_image.copy()
        overlay[mask] = (0, 255, 0)  # 将分割区域设为绿色
        cv2.addWeighted(overlay, 0.5, segmented_image, 0.5, 0, segmented_image)

        # 保存分割图像供窗口调整时使用
        self.segmented_image = segmented_image

        # 显示分割结果
        self.display_segmented_image(segmented_image)
        self.display_single_object(mask)

        self.statusbar.showMessage("SAM分割完成")
    except Exception as e:
        QMessageBox.critical(self, "错误", f"SAM分割失败: {str(e)}")
        self.statusbar.showMessage("SAM分割失败")
    finally:
        # 清理选点数据和历史
        self.input_points = []
        self.input_labels = []
        self.input_points_history = []
        self.input_labels_history = []
        self.drawing_image_history = []
        self.sam_mode = False
        # 恢复原图显示
        self.display_image()
